import json
import csv
import logging
import os
import re
from collections import deque

from sonetto.agents.base import BaseAgent, ContextLevel
from sonetto.agents.chat import ChatAgent
from sonetto.agents.tool import ToolAgent
from sonetto.core.message import Message
from sonetto.core.routing.broadcast import CommonBroadcastConditions
from sonetto.core.scheduler import Scheduler
from sonetto.core.toolset import McpToolset, field

base_url_glm = "https://open.bigmodel.cn/api/paas/v4"
api_key_glm = "ff37fbba1e80d4cb5f38ac819893f4a1.c6T0Zm3Sd9YODRLV"
model_glm = "glm-4-flash"

base_url_ds = "https://api.deepseek.com"
api_key_ds = "sk-3bb4d61950df4f4d833f800848e1f461"
model_ds = "deepseek-chat"

toolset = McpToolset()


class CoordinatorAgent(BaseAgent):

    def __init__(self, agent_id: str, working_directory: str):
        super().__init__(agent_id)
        self.directory = working_directory
        self.document_queue = deque([
            doc for doc in os.listdir(working_directory) if doc.endswith('md')
        ])
        logging.debug(f"文档队列：{self.document_queue}")
        self.current_document = None
        self.current_output = None
        self.current_line = 0
        self.total_lines = 0
        self.window_size = 100
        self.state = 'chunk'
        self.bounds, self.chunks = [], deque()

    def list_directory(self):
        return os.listdir(self.directory)

    def pretty_print(self, content: dict[str], lineno: bool = True) -> str:
        ret = ""
        for key, value in content.items():
            if lineno:
                ret += f"第{key}行: {value}\n"
            else:
                ret += f"{value}\n"
        ret += "-----------------\n"
        return ret

    def read_lines(self, start: int, end: int, lineno: bool = True) -> str:
        """
        读取文件中指定范围的行
        
        Args:
            filename: 文件路径
            start: 起始行号（从1开始，包含）
            end: 结束行号（从1开始，包含）
            
        Returns:
            包含行号和内容的字符串列表，每个字符串格式为 "行号: 内容"
            
        Raises:
            FileNotFoundError: 文件不存在时抛出
            ValueError: 起始或结束行号无效时抛出
            TypeError: 行号不是整数时抛出
        """
        # 参数验证
        if not isinstance(start, int) or not isinstance(end, int):
            raise TypeError("行号必须是整数")

        if start < 1:
            raise ValueError("起始行号必须大于等于1")

        if end < start:
            raise ValueError("结束行号不能小于起始行号")

        try:
            result = {}
            with open(self.current_document, 'r', encoding='utf-8') as file:
                line_num = 0
                for line in file:
                    line_num += 1

                    if line_num < start:
                        continue

                    if line_num <= end:
                        result[line_num] = f"{line.rstrip()}"
                    else:
                        break

                if not result and start > line_num:
                    raise ValueError(f"文件只有{line_num}行，少于起始行号{start}")
                self.current_line += self.window_size
                return self.pretty_print(result, lineno)

        except FileNotFoundError:
            raise FileNotFoundError(f"文件 '{self.current_document}' 不存在")

    def split_md2jsonl(self) -> None:
        self.bounds.insert(0, 1)
        self.bounds.append(self.total_lines + 1)
        print(self.bounds)
        with open(self.current_document, 'r', encoding='utf-8') as file:
            with open(self.current_output, 'w', encoding='utf-8') as output:
                writer = csv.writer(output, quoting=csv.QUOTE_MINIMAL, escapechar='\\')
                for i in range(1, len(self.bounds)):
                    content = self.read_lines(self.bounds[i - 1],
                                              self.bounds[i] - 1, False)
                    # output.write(
                    #     json.dumps({'content': content}, ensure_ascii=False) + '\n')
                    # writer.writerow([json.dumps(content, ensure_ascii=False)])
                    writer.writerow([content])
                    self.chunks.append(content)
        self.bounds.clear()

    def count_lines(self) -> int:
        try:
            with open(self.current_document, 'r', encoding='utf-8') as file:
                return len(file.readlines())
        except FileNotFoundError:
            raise FileNotFoundError(f"文件 '{self.current_document}' 不存在")

    def turn2next_document(self) -> None:
        if self.document_queue:
            self.current_document = os.path.join(self.directory,
                                                 self.document_queue.popleft())
            self.current_output = f"{self.current_document}.csv"
            self.detail_output = f"{self.current_document}_propositions.csv"
            self.current_line = 1
            self.total_lines = self.count_lines()
            return True
        else:
            return False

    def parse_bounds(self, content: str) -> list[int]:
        bounds = re.findall(r"```json\n(.+?)\n```", content, re.DOTALL)
        print(bounds)
        bounds = json.loads(bounds[0])
        self.bounds.extend([
            ele for ele in bounds
            if ele % self.window_size != 0 and ele % self.window_size != 1
        ])

    def parse_propositions(self, content: str) -> list[str]:
        propositions = re.findall(r"```json\n(.+?)\n```", content, re.DOTALL)
        propositions = json.loads(propositions[0])
        with open(self.detail_output, 'a', encoding='utf-8') as output:
            writer = csv.writer(output)
            for prop in propositions:
                writer.writerow([prop])
        return propositions

    def next_chunk(self) -> str:
        if self.chunks:
            return self.chunks.popleft()
        else:
            return ""

    def forward(self, message: Message) -> Message:
        sender = message.sender
        receiver = message.receiver
        content = message.content

        if sender == '__init__':
            if not self.turn2next_document():
                return Message(content="没有文档需要处理",
                               sender=self.agent_id,
                               receiver='chunker')
            else:
                content = f"""
接下来开始对文档 '{self.current_document}' 进行分块。首先需要分块的文本内容如下所示：
```markdown
{self.read_lines(self.current_line, self.current_line + self.window_size - 1)}
```
请确保每个分块主题的独立完整性，并严格遵照既定格式回复。
"""
                return Message(content=content,
                               sender=self.agent_id,
                               receiver='chunker')

        # 检查是否有文档需要处理
        if self.current_line >= self.total_lines and self.state == 'chunk':
            '''
```json
[5, 11, 29, 33, 39, 41, 42, 99]
```
            '''
            self.parse_bounds(content)
            self.split_md2jsonl()
            self.state = 'split'
            content = f"""
请将以下文本分解为命题：
```markdown
{self.next_chunk()}
```
"""
            return Message(content=content,
                           sender=self.agent_id,
                           receiver='spliter')
        if sender == 'chunker':
            # 提取json数组中的元素，拼接到self.bounds中
            self.parse_bounds(content)
            content = f"""
继续对文档 '{self.current_document}' 进行分块。接下来需要分块的文本内容如下所示：
```markdown
{self.read_lines(self.current_line, self.current_line + self.window_size - 1)}
```
请确保每个分块主题的独立完整性，并严格遵照既定格式回复。
"""
            return Message(content=content,
                           sender=self.agent_id,
                           receiver=sender)
        elif sender == 'spliter':
            self.parse_propositions(content)
            chunk = self.next_chunk()
            if chunk:
                content = f"""
请将以下文本分解为命题：
```markdown
{chunk}
```
"""
                return Message(content=content,
                               sender=self.agent_id,
                               receiver=sender)
            else:
                self.parse_propositions(content)
                self.state = 'chunk'
                if not self.turn2next_document():
                    return Message(content="所有文档均已整理完毕，圆满结束！！！",
                                   sender=self.agent_id,
                                   receiver=sender)
                else:
                    content = f"""
    上一个文档已完成分块，接下来开始对文档 '{self.current_document}' 进行分块。首先需要分块的文本内容如下所示：
    ```markdown
    {self.read_lines(self.current_line, self.current_line + self.window_size - 1)}
    ```
    请确保每个分块主题的独立完整性，并严格遵照既定格式回复。
    """
                    return Message(content=content,
                                   sender=self.agent_id,
                                   receiver='chunker')


coordinator = CoordinatorAgent(agent_id="coordinator",
                               working_directory=".\HSBC")

chunker = ChatAgent(agent_id="chunker",
                    base_url=base_url_ds,
                    api_key=api_key_ds,
                    model=model_ds,
                    system_prompt='''
你是一个文档语义边界划分专家，你的任务是识别文本的语义边界。
你划分的块最后会作为问答系统的参考内容，故其语义的独立完整性非常重要。

### 语义边界划分指南：
1. 语义边界之间的块应在主题/语义上连贯且独立。
2. 语义边界划分：
    - 分节边界（章节、标题）
    - 内容主题转换
    - 形成逻辑单元的段落组
    - 问答对（先问再答）
3. 图片或表格不要单独划分，应当和最近的文本块合并。
4. 分块粒度不宜太粗（比如分成整个章节），也不宜太细（比如按句划分）。
5. 边界行的上面部分（不含边界行）会被归为上一个块，下面部分（包含边界行）会被归为下一个块。

由于文档非常庞大，用户只能分多次为你提供文档内容，每次只能提供100行。

请你具体分析语义边界的划分原因，最后以json数组的格式给出你的语义边界行号，例：
/*此处分析每个分块的主题和语义独立完整性……*/
```json
[5, 36, 54]
```
''')

spliter = ChatAgent(agent_id="spliter",
                    base_url=base_url_ds,
                    api_key=api_key_ds,
                    model=model_ds,
                    system_prompt="""
将文本内容分解为清晰简单的命题，确保它们脱离上下文也能被理解。这些命题将用于问答系统的知识库中，因此它们应该是独立的、完整的，并且不依赖于上下文。

1. 将复合句拆分为简单句。尽可能保留输入中的原始表述，但也确保每个命题至少包含主谓宾。

2. 对于任何伴随有额外描述信息的命名实体，将这些信息分离成独立的命题。

3. 通过给名词或整个句子添加必要的修饰语，以及将代词（如"它"、"他"、"她"、"他们"、"这个"、"那个"）替换为它们所指代实体的全名，使命题去语境化。

4. 去除缺少有效信息的命题，如目录、广告词以及其他不可能作为某个问题答案的命题。

5. 将结果以JSON格式的字符串列表呈现。

输入: 
```markdown
关于复活节兔(Osterhase)的最早证据是由医学教授乔治·弗兰克·冯·弗兰肯瑙于1678年在德国西南部记录的，但直到18世纪它在德国其他地区仍不为人所知。学者理查德·瑟蒙写道："兔子经常在春天的花园中被发现，因此可能成为解释隐藏在那里供儿童寻找的彩蛋起源的一个便利解释。或者，有一种欧洲传统认为兔子会下蛋，因为兔子的窝和田凫鸟的巢看起来非常相似，两者都出现在草地上，并且都是在春天首次被发现。在十九世纪，复活节贺卡、玩具和书籍的影响使得复活节兔在整个欧洲变得流行。德国移民随后将这一习俗传播到英国和美国，在那里它演变成了复活节兔。"
```

输出: 
```json
[ "关于复活节兔的最早证据是由乔治·弗兰克·冯·弗兰肯瑙于1678年在德国西南部记录的。", "乔治·弗兰克·冯·弗兰肯瑙是一位医学教授。", "复活节兔的证据直到18世纪在德国其他地区仍不为人所知。", "理查德·瑟蒙是一位学者。", "理查德·瑟蒙写出了一个关于兔子与复活节传统之间联系的可能解释的假设。", "兔子经常在春天的花园中被发现。", "兔子可能成为解释隐藏在花园中供儿童寻找的彩蛋起源的一个便利解释。", "有一种欧洲传统认为兔子会下蛋。", "兔子的窝和田凫鸟的巢看起来非常相似。", "兔子和田凫鸟的巢都出现在草地上，并且都是在春天首次被发现。", "在十九世纪，复活节贺卡、玩具和书籍的影响使得复活节兔在整个欧洲变得流行。", "德国移民将复活节兔的习俗传播到英国和美国。", "复活节兔的习俗在英国和美国演变成了复活节兔。" ]
```
""")

scheduler = Scheduler()
scheduler.register_agent(coordinator)
scheduler.register_agent(chunker)
scheduler.register_agent(spliter)
scheduler.set_route(chunker.agent_id, coordinator.agent_id)
scheduler.set_route(spliter.agent_id, coordinator.agent_id)

if __name__ == "__main__":
    logging.debug("文档分块任务开始")
    scheduler.run_stream("请完成文档分块任务",
                         cancel_token="圆满结束",
                         cancel_from=[coordinator.agent_id],
                         max_step=1500)
    logging.debug("文档分块任务结束")
